import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * Creat with IntelliJ IDEA
 * Description:
 * User:WZW
 * Date:2022-02-25
 * Time:15:20
 */
public class TopK {
    /*
    求数组中前K个最小的元素
     */
    public static int[] topK(int[] array, int k) {
        //1.创建一个大小为k的大根堆
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });
        //2.遍历数组当中个元素，前k个元素放到队列当中
        for (int i = 0; i < array.length; i++) {
            if (maxHeap.size() < k) {
                maxHeap.offer(array[i]);
            } else {
                //3.从第k+1个元素开始，每个元素和堆顶元素进行比较
                int top = maxHeap.peek();
                if (top > array[i]) {
                    //1.先弹出
                    maxHeap.poll();
                    //2.后存入
                    maxHeap.offer(array[i]);
                }
            }
        }
        int[] tmp = new int[k];
        for (int i = 0; i < k; i++) {
            tmp[i] = maxHeap.poll();
        }
        return tmp;
    }

    /*
    查找和最小的 K 对数字
    https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/
     */
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<>
            (k, new Comparator<List<Integer>>() {
                public int compare(List<Integer> o1, List<Integer> o2) {
                    return (o2.get(0) + o2.get(1)) - (o1.get(0) + o1.get(1));
                }
            });
        for(int i = 0; i<Math.min(nums1.length,k);i++) {
            for (int j = 0; j < Math.min(nums2.length, k); j++) {
                if (maxHeap.size() < k) {
                    List<Integer> tmp = new ArrayList<>();
                    tmp.add(nums1[i]);
                    tmp.add(nums2[j]);
                    maxHeap.offer(tmp);
                } else {
                    List<Integer> top = maxHeap.peek();
                    if (top.get(0) + top.get(1) > nums1[i] + nums2[j]) {
                        maxHeap.poll();
                        List<Integer> tmp = new ArrayList<>();
                        tmp.add(nums1[i]);
                        tmp.add(nums2[j]);
                        maxHeap.offer(tmp);
                    }
                }
            }
        }
        List<List<Integer>> ret = new ArrayList<>();
            for(int i = 0;i<k &&!maxHeap.isEmpty();i++){
            ret.add(maxHeap.poll());
        }
            return ret;
    }

    public List<List<Integer>> kSmallestPairs1(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return (nums1[o1[0]]+nums2[o1[1]])-(nums1[o2[0]]+nums2[o2[1]]);
            }
        });
        for (int i = 0; i < Math.min(nums1.length,k); i++) {
            pq.offer(new int[]{i,0});
        }
        List<List<Integer>> ret = new ArrayList<>();
        for (int i = 0; i < k && !pq.isEmpty(); i++) {
            int[] index = pq.poll();
            List<Integer> tmpList = new ArrayList<>();
            tmpList.add(nums1[index[0]]);
            tmpList.add(nums2[index[1]]);
            ret.add(tmpList);
            if(index[1]+1 < nums2.length){
                pq.offer(new int[]{index[0],index[1]+1});
            }
        }
        return ret;
    }
}
